// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <melon/hash.h>
#include <melon/container/f14_map.h>

#include <pollux/common/base/exceptions.h>
#include <pollux/common/base/runtime_metrics.h>
#include <pollux/common/exception/exception.h>

namespace kumo::pollux::dwio::common {
    // Common base for writer version information used in interpreting
    // metadata. Needed to have format-independent signatures for
    // format-specific functions. Each format implementation downcasts this to the
    // format-specific metadata.
    struct StatsContext {
        virtual ~StatsContext() = default;
    };

    struct KeyInfo {
    public:
        explicit KeyInfo(int64_t intKey)
            : intKey{std::make_optional<int64_t>(intKey)} {
        }

        explicit KeyInfo(const std::string &bytesKey)
            : bytesKey{std::make_optional<std::string>(bytesKey)} {
        }

        bool operator==(const KeyInfo &other) const {
            return intKey == other.intKey && bytesKey == other.bytesKey;
        }

        std::string toString() const {
            if (intKey.has_value()) {
                return melon::to<std::string>(*intKey);
            } else if (bytesKey.has_value()) {
                return *bytesKey;
            }
            POLLUX_UNREACHABLE("Illegal null key info");
        }

        std::optional<int64_t> intKey;
        std::optional<std::string> bytesKey;

    private:
        KeyInfo() {
        }
    };

    struct KeyInfoHash {
        KeyInfoHash() = default;

        size_t operator()(const KeyInfo &keyInfo) const {
            if (keyInfo.intKey.has_value()) {
                return melon::Hash{}(*keyInfo.intKey);
            } else if (keyInfo.bytesKey.has_value()) {
                return melon::Hash{}(*keyInfo.bytesKey);
            }
            POLLUX_UNREACHABLE("Illegal null key info");
        }
    };

    /**
     * Statistics that are available for all types of columns.
     */
    class ColumnStatistics {
    public:
        ColumnStatistics(
            std::optional<uint64_t> valueCount,
            std::optional<bool> hasNull,
            std::optional<uint64_t> rawSize,
            std::optional<uint64_t> size,
            std::optional<int64_t> numDistinct = std::nullopt)
            : valueCount_(valueCount),
              hasNull_(hasNull),
              rawSize_(rawSize),
              size_(size),
              numDistinct_(numDistinct) {
        }

        virtual ~ColumnStatistics() = default;

        /**
         * Get the number of values in this column. It will differ from the number
         * of rows because of NULL values and repeated (list/map) values.
         */
        std::optional<uint64_t> getNumberOfValues() const {
            return valueCount_;
        }

        /**
         * Get whether column has null value.
         *
         * WARNING: Some writer implementation does not take ancestor nulls into
         * account, so this value should not be trusted.  Check whether
         * `getNumberOfValues()' is smaller than the row group size is a more accurate
         * way.
         */
        std::optional<bool> hasNull() const {
            return hasNull_;
        }

        /**
         * Get uncompressed size of all data including child
         */
        std::optional<uint64_t> getRawSize() const {
            return rawSize_;
        }

        /**
         * Get total length of all streams including child.
         */
        std::optional<uint64_t> getSize() const {
            return size_;
        }

        std::optional<uint64_t> numDistinct() const {
            return numDistinct_;
        }

        void setNumDistinct(int64_t count) {
            POLLUX_CHECK(
                !numDistinct_.has_value(), "numDistinct_ can be set only once.");
            numDistinct_ = count;
        }

        /**
         * return string representation of this stats object
         */
        virtual std::string toString() const {
            return melon::to<std::string>(
                "RawSize: ",
                (rawSize_ ? melon::to<std::string>(rawSize_.value()) : "unknown"),
                ", Size: ",
                (size_ ? melon::to<std::string>(size_.value()) : "unknown"),
                ", Values: ",
                (valueCount_ ? melon::to<std::string>(valueCount_.value()) : "unknown"),
                ", hasNull: ",
                (hasNull_ ? (hasNull_.value() ? "yes" : "no") : "unknown"));
        }

    protected:
        ColumnStatistics() {
        }

        std::optional<uint64_t> valueCount_;
        std::optional<bool> hasNull_;
        std::optional<uint64_t> rawSize_;
        std::optional<uint64_t> size_;
        std::optional<uint64_t> numDistinct_;
    };

    /**
     * Statistics for binary columns.
     */
    class BinaryColumnStatistics : public virtual ColumnStatistics {
    public:
        BinaryColumnStatistics(
            std::optional<uint64_t> valueCount,
            std::optional<bool> hasNull,
            std::optional<uint64_t> rawSize,
            std::optional<uint64_t> size,
            std::optional<uint64_t> length)
            : ColumnStatistics(valueCount, hasNull, rawSize, size), length_(length) {
        }

        BinaryColumnStatistics(
            const ColumnStatistics &colStats,
            std::optional<uint64_t> length)
            : ColumnStatistics(colStats), length_(length) {
        }

        ~BinaryColumnStatistics() override = default;

        /**
         * get optional total length
         */
        std::optional<uint64_t> getTotalLength() const {
            return length_;
        }

        std::string toString() const override {
            return melon::to<std::string>(
                ColumnStatistics::toString(),
                ", Length: ",
                (length_.has_value()
                     ? melon::to<std::string>(length_.value())
                     : "unknown"));
        }

    protected:
        BinaryColumnStatistics() {
        }

        std::optional<uint64_t> length_;
    };

    /**
     * Statistics for boolean columns.
     */
    class BooleanColumnStatistics : public virtual ColumnStatistics {
    public:
        BooleanColumnStatistics(
            std::optional<uint64_t> valueCount,
            std::optional<bool> hasNull,
            std::optional<uint64_t> rawSize,
            std::optional<uint64_t> size,
            std::optional<uint64_t> trueCount)
            : ColumnStatistics(valueCount, hasNull, rawSize, size),
              trueCount_(trueCount) {
        }

        BooleanColumnStatistics(
            const ColumnStatistics &colStats,
            std::optional<uint64_t> trueCount)
            : ColumnStatistics(colStats), trueCount_(trueCount) {
        }

        ~BooleanColumnStatistics() override = default;

        /*
         * get optional true count
         */
        std::optional<uint64_t> getTrueCount() const {
            return trueCount_;
        }

        /*
         * get optional false count
         */
        std::optional<uint64_t> getFalseCount() const {
            auto valueCount = getNumberOfValues();
            return trueCount_.has_value() && valueCount.has_value()
                       ? valueCount.value() - trueCount_.value()
                       : std::optional<uint64_t>();
        }

        std::string toString() const override {
            return melon::to<std::string>(
                ColumnStatistics::toString(),
                ", trueCount: ",
                (trueCount_.has_value()
                     ? melon::to<std::string>(trueCount_.value())
                     : "unknown"));
        }

    protected:
        BooleanColumnStatistics() {
        }

        std::optional<uint64_t> trueCount_;
    };

    /**
     * Statistics for float and double columns.
     */
    class DoubleColumnStatistics : public virtual ColumnStatistics {
    public:
        DoubleColumnStatistics(
            std::optional<uint64_t> valueCount,
            std::optional<bool> hasNull,
            std::optional<uint64_t> rawSize,
            std::optional<uint64_t> size,
            std::optional<double> min,
            std::optional<double> max,
            std::optional<double> sum)
            : ColumnStatistics(valueCount, hasNull, rawSize, size),
              min_(min),
              max_(max),
              sum_(sum) {
        }

        DoubleColumnStatistics(
            const ColumnStatistics &colStats,
            std::optional<double> min,
            std::optional<double> max,
            std::optional<double> sum)
            : ColumnStatistics(colStats), min_(min), max_(max), sum_(sum) {
        }

        ~DoubleColumnStatistics() override = default;

        /**
         * Get optional smallest value in the column. Only defined if
         * getNumberOfValues is non-zero.
         */
        std::optional<double> getMinimum() const {
            return min_;
        }

        /**
         * Get optional largest value in the column. Only defined if getNumberOfValues
         * is non-zero.
         */
        std::optional<double> getMaximum() const {
            return max_;
        }

        /**
         * Get optional sum of the values in the column.
         */
        std::optional<double> getSum() const {
            return sum_;
        }

        std::string toString() const override {
            return melon::to<std::string>(
                ColumnStatistics::toString(),
                ", min: ",
                (min_.has_value() ? melon::to<std::string>(min_.value()) : "unknown"),
                ", max: ",
                (max_.has_value() ? melon::to<std::string>(max_.value()) : "unknown"),
                ", sum: ",
                (sum_.has_value() ? melon::to<std::string>(sum_.value()) : "unknown"));
        }

    protected:
        DoubleColumnStatistics() {
        }

        std::optional<double> min_;
        std::optional<double> max_;
        std::optional<double> sum_;
    };

    /**
     * Statistics for all of the integer columns, such as byte, short, int, and
     * long.
     */
    class IntegerColumnStatistics : public virtual ColumnStatistics {
    public:
        IntegerColumnStatistics(
            std::optional<uint64_t> valueCount,
            std::optional<bool> hasNull,
            std::optional<uint64_t> rawSize,
            std::optional<uint64_t> size,
            std::optional<int64_t> min,
            std::optional<int64_t> max,
            std::optional<int64_t> sum)
            : ColumnStatistics(valueCount, hasNull, rawSize, size),
              min_(min),
              max_(max),
              sum_(sum) {
        }

        IntegerColumnStatistics(
            const ColumnStatistics &colStats,
            std::optional<int64_t> min,
            std::optional<int64_t> max,
            std::optional<int64_t> sum)
            : ColumnStatistics(colStats), min_(min), max_(max), sum_(sum) {
        }

        ~IntegerColumnStatistics() override = default;

        /**
         * Get optional smallest value in the column. Only defined if
         * getNumberOfValues is non-zero.
         */
        std::optional<int64_t> getMinimum() const {
            return min_;
        }

        /**
         * Get optional largest value in the column. Only defined if getNumberOfValues
         * is non-zero.
         */
        std::optional<int64_t> getMaximum() const {
            return max_;
        }

        /**
         * Get optional sum of the column. Only valid if getNumberOfValues is non-zero
         * and sum doesn't overflow
         */
        std::optional<int64_t> getSum() const {
            return sum_;
        }

        std::string toString() const override {
            return melon::to<std::string>(
                ColumnStatistics::toString(),
                ", min: ",
                (min_.has_value() ? melon::to<std::string>(min_.value()) : "unknown"),
                ", max: ",
                (max_.has_value() ? melon::to<std::string>(max_.value()) : "unknown"),
                ", sum: ",
                (sum_.has_value() ? melon::to<std::string>(sum_.value()) : "unknown"));
        }

    protected:
        IntegerColumnStatistics() {
        }

        std::optional<int64_t> min_;
        std::optional<int64_t> max_;
        std::optional<int64_t> sum_;
    };

    /**
     * Statistics for string columns.
     */
    class StringColumnStatistics : public virtual ColumnStatistics {
    public:
        StringColumnStatistics(
            std::optional<uint64_t> valueCount,
            std::optional<bool> hasNull,
            std::optional<uint64_t> rawSize,
            std::optional<uint64_t> size,
            std::optional<std::string> min,
            std::optional<std::string> max,
            std::optional<int64_t> length)
            : ColumnStatistics(valueCount, hasNull, rawSize, size),
              min_(min),
              max_(max),
              length_(length) {
        }

        StringColumnStatistics(
            const ColumnStatistics &colStats,
            std::optional<std::string> min,
            std::optional<std::string> max,
            std::optional<int64_t> length)
            : ColumnStatistics(colStats), min_(min), max_(max), length_(length) {
        }

        ~StringColumnStatistics() override = default;

        /**
         * Get optional minimum value for the column.
         */
        const std::optional<std::string> &getMinimum() const {
            return min_;
        }

        /**
         * Get optional maximum value for the column.
         */
        const std::optional<std::string> &getMaximum() const {
            return max_;
        }

        /**
         * Get optional total length of all values.
         */
        std::optional<uint64_t> getTotalLength() const {
            return length_;
        }

        std::string toString() const override {
            return melon::to<std::string>(
                ColumnStatistics::toString(),
                ", min: ",
                min_.value_or("unknown"),
                ", max: ",
                max_.value_or("unknown"),
                ", length: ",
                (length_.has_value()
                     ? melon::to<std::string>(length_.value())
                     : "unknown"));
        }

    protected:
        StringColumnStatistics() {
        }

        std::optional<std::string> min_;
        std::optional<std::string> max_;
        std::optional<uint64_t> length_;
    };

    /**
     * Statistics for (flat) map columns.
     */
    class MapColumnStatistics : public virtual ColumnStatistics {
    public:
        MapColumnStatistics(
            std::optional<uint64_t> valueCount,
            std::optional<bool> hasNull,
            std::optional<uint64_t> rawSize,
            std::optional<uint64_t> size,
            melon::F14FastMap<
                KeyInfo,
                std::unique_ptr<ColumnStatistics>,
                melon::transparent<KeyInfoHash> > &&entryStatistics)
            : ColumnStatistics(valueCount, hasNull, rawSize, size),
              entryStatistics_{std::move(entryStatistics)} {
        }

        ~MapColumnStatistics() override = default;

        const melon::F14FastMap<
            KeyInfo,
            std::unique_ptr<ColumnStatistics>,
            melon::transparent<KeyInfoHash> > &
        getEntryStatistics() const {
            return entryStatistics_;
        }

        std::string toString() const override {
            std::vector<std::string> values{};
            values.reserve(entryStatistics_.size());
            for (const auto &entry: entryStatistics_) {
                auto &stats = *entry.second;
                values.push_back(fmt::format(
                    "{{ Key: {}, Stats: {},}}",
                    entry.first.toString(),
                    stats.toString()));
            }
            std::string repr;
            melon::join(",", values, repr);
            return melon::to<std::string>(ColumnStatistics::toString(), repr);
        }

    protected:
        MapColumnStatistics()
            : entryStatistics_{17, melon::transparent<KeyInfoHash>()} {
        }

        melon::F14FastMap<
            KeyInfo,
            std::unique_ptr<ColumnStatistics>,
            melon::transparent<KeyInfoHash> >
        entryStatistics_;
    };

    class Statistics {
    public:
        virtual ~Statistics() = default;

        /**
         * Get the statistics of the given column.
         * @param colId id of the column
         * @return one column's statistics
         */
        virtual const ColumnStatistics &getColumnStatistics(uint32_t colId) const = 0;

        /**
         * Get the number of columns
         * @return the number of columns
         */
        virtual uint32_t getNumberOfColumns() const = 0;
    };

    struct ColumnReaderStatistics {
        // Number of rows returned by string dictionary reader that is flattened
        // instead of keeping dictionary encoding.
        int64_t flattenStringDictionaryValues{0};
    };

    struct RuntimeStatistics {
        // Number of splits skipped based on statistics.
        int64_t skippedSplits{0};

        // Number of splits processed based on statistics.
        int64_t processedSplits{0};

        // Total bytes in splits skipped based on statistics.
        int64_t skippedSplitBytes{0};

        // Number of strides (row groups) skipped based on statistics.
        int64_t skippedStrides{0};

        // Number of strides (row groups) processed based on statistics.
        int64_t processedStrides{0};

        int64_t footerBufferOverread{0};

        int64_t numStripes{0};

        ColumnReaderStatistics columnReaderStatistics;

        std::unordered_map<std::string, RuntimeCounter> toMap() {
            std::unordered_map<std::string, RuntimeCounter> result;
            if (skippedSplits > 0) {
                result.emplace("skippedSplits", RuntimeCounter(skippedSplits));
            }
            if (processedSplits > 0) {
                result.emplace("processedSplits", RuntimeCounter(processedSplits));
            }
            if (skippedSplitBytes > 0) {
                result.emplace(
                    "skippedSplitBytes",
                    RuntimeCounter(skippedSplitBytes, RuntimeCounter::Unit::kBytes));
            }
            if (skippedStrides > 0) {
                result.emplace("skippedStrides", RuntimeCounter(skippedStrides));
            }
            if (processedStrides > 0) {
                result.emplace("processedStrides", RuntimeCounter(processedStrides));
            }
            if (footerBufferOverread > 0) {
                result.emplace(
                    "footerBufferOverread",
                    RuntimeCounter(footerBufferOverread, RuntimeCounter::Unit::kBytes));
            }
            if (numStripes > 0) {
                result.emplace("numStripes", RuntimeCounter(numStripes));
            }
            if (columnReaderStatistics.flattenStringDictionaryValues > 0) {
                result.emplace(
                    "flattenStringDictionaryValues",
                    RuntimeCounter(columnReaderStatistics.flattenStringDictionaryValues));
            }
            return result;
        }
    };
} // namespace kumo::pollux::dwio::common
